// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/compute/api_aggregate.h>
#include <nebula/compute/kernels/aggregate_internal.h>
#include <nebula/compute/kernels/common_internal.h>
#include <nebula/bits/bit_run_reader.h>
#include <turbo/utility/tdigest.h>

namespace nebula::compute::internal {

    namespace {
        using nebula::internal::VisitSetBitRunsVoid;

        template<typename ArrowType>
        struct TDigestImpl : public ScalarAggregator {
            using ThisType = TDigestImpl<ArrowType>;
            using ArrayType = typename TypeTraits<ArrowType>::ArrayType;
            using CType = typename TypeTraits<ArrowType>::CType;

            TDigestImpl(const TDigestOptions &options, const DataType &in_type)
                    : options{options},
                      tdigest{options.delta, options.buffer_size},
                      count{0},
                      decimal_scale{0},
                      all_valid{true} {
                if (is_decimal_type<ArrowType>::value) {
                    decimal_scale = turbo::checked_cast<const DecimalType &>(in_type).scale();
                }
            }

            template<typename T>
            double ToDouble(T value) const {
                return static_cast<double>(value);
            }

            double ToDouble(const Decimal128 &value) const { return value.ToDouble(decimal_scale); }

            double ToDouble(const Decimal256 &value) const { return value.ToDouble(decimal_scale); }

            turbo::Status consume(KernelContext *, const ExecSpan &batch) override {
                if (!this->all_valid) return turbo::OkStatus();
                if (!options.skip_nulls && batch[0].null_count() > 0) {
                    this->all_valid = false;
                    return turbo::OkStatus();
                }
                if (batch[0].is_array()) {
                    const ArraySpan &data = batch[0].array;
                    const CType *values = data.get_values<CType>(1);

                    if (data.length > data.get_null_count()) {
                        this->count += data.length - data.get_null_count();
                        VisitSetBitRunsVoid(data.buffers[0].data, data.offset, data.length,
                                            [&](int64_t pos, int64_t len) {
                                                for (int64_t i = 0; i < len; ++i) {
                                                    this->tdigest.nan_add(ToDouble(values[pos + i]));
                                                }
                                            });
                    }
                } else {
                    const CType value = UnboxScalar<ArrowType>::Unbox(*batch[0].scalar);
                    if (batch[0].scalar->is_valid) {
                        this->count += 1;
                        for (int64_t i = 0; i < batch.length; i++) {
                            this->tdigest.nan_add(ToDouble(value));
                        }
                    }
                }
                return turbo::OkStatus();
            }

            turbo::Status MergeFrom(KernelContext *, KernelState &&src) override {
                const auto &other = turbo::checked_cast<const ThisType &>(src);
                if (!this->all_valid || !other.all_valid) {
                    this->all_valid = false;
                    return turbo::OkStatus();
                }
                this->tdigest.merge(other.tdigest);
                this->count += other.count;
                return turbo::OkStatus();
            }

            turbo::Status Finalize(KernelContext *ctx, Datum *out) override {
                const int64_t out_length = options.q.size();
                auto out_data = ArrayData::create(float64(), out_length, 0);
                out_data->buffers.resize(2, nullptr);
                TURBO_MOVE_OR_RAISE(out_data->buffers[1],
                                       ctx->allocate(out_length * sizeof(double)));
                double *out_buffer = out_data->template GetMutableValues<double>(1);

                if (this->tdigest.is_empty() || !this->all_valid || this->count < options.min_count) {
                    TURBO_MOVE_OR_RAISE(out_data->buffers[0], ctx->allocate_bitmap(out_length));
                    std::memset(out_data->buffers[0]->mutable_data(), 0x00,
                                out_data->buffers[0]->size());
                    std::fill(out_buffer, out_buffer + out_length, 0.0);
                    out_data->null_count = out_length;
                } else {
                    for (int64_t i = 0; i < out_length; ++i) {
                        out_buffer[i] = this->tdigest.quantile(this->options.q[i]);
                    }
                }
                *out = Datum(std::move(out_data));
                return turbo::OkStatus();
            }

            const TDigestOptions options;
            turbo::TDigest tdigest;
            int64_t count;
            int32_t decimal_scale;
            bool all_valid;
        };

        struct TDigestInitState {
            std::unique_ptr<KernelState> state;
            KernelContext *ctx;
            const DataType &in_type;
            const TDigestOptions &options;

            TDigestInitState(KernelContext *ctx, const DataType &in_type,
                             const TDigestOptions &options)
                    : ctx(ctx), in_type(in_type), options(options) {}

            turbo::Status Visit(const DataType &) {
                return turbo::unimplemented_error("No tdigest implemented");
            }

            turbo::Status Visit(const Fp16Type &) {
                return turbo::unimplemented_error("No tdigest implemented");
            }

            template<typename Type>
            enable_if_number<Type, turbo::Status> Visit(const Type &) {
                state.reset(new TDigestImpl<Type>(options, in_type));
                return turbo::OkStatus();
            }

            template<typename Type>
            enable_if_decimal<Type, turbo::Status> Visit(const Type &) {
                state.reset(new TDigestImpl<Type>(options, in_type));
                return turbo::OkStatus();
            }

            turbo::Result<std::unique_ptr<KernelState>> create() {
                TURBO_RETURN_NOT_OK(visit_type_inline(in_type, this));
                return std::move(state);
            }
        };

        turbo::Result<std::unique_ptr<KernelState>> TDigestInit(KernelContext *ctx,
                                                         const KernelInitArgs &args) {
            TDigestInitState visitor(ctx, *args.inputs[0].type,
                                     static_cast<const TDigestOptions &>(*args.options));
            return visitor.create();
        }

        void AddTDigestKernels(KernelInit init,
                               const std::vector<std::shared_ptr<DataType>> &types,
                               ScalarAggregateFunction *func) {
            for (const auto &ty: types) {
                auto sig = KernelSignature::create({InputType(ty->id())}, float64());
                add_agg_kernel(std::move(sig), init, func);
            }
        }

        const FunctionDoc tdigest_doc{
                "Approximate quantiles of a numeric array with T-Digest algorithm",
                ("By default, 0.5 quantile (median) is returned.\n"
                 "Nulls and NaNs are ignored.\n"
                 "An array of nulls is returned if there is no valid data point."),
                {"array"},
                "TDigestOptions"};

        const FunctionDoc approximate_median_doc{
                "Approximate median of a numeric array with T-Digest algorithm",
                ("Nulls and NaNs are ignored.\n"
                 "A null scalar is returned if there is no valid data point."),
                {"array"},
                "ScalarAggregateOptions"};

        std::shared_ptr<ScalarAggregateFunction> add_tdigest_agg_kernels() {
            static auto default_tdigest_options = TDigestOptions::defaults();
            auto func = std::make_shared<ScalarAggregateFunction>(
                    "tdigest", Arity::Unary(), tdigest_doc, &default_tdigest_options);
            AddTDigestKernels(TDigestInit, numeric_types(), func.get());
            AddTDigestKernels(TDigestInit, {decimal128(1, 1), decimal256(1, 1)}, func.get());
            return func;
        }

        std::shared_ptr<ScalarAggregateFunction> add_approximate_median_agg_kernels(
                const ScalarAggregateFunction *tdigest_func) {
            static ScalarAggregateOptions default_scalar_aggregate_options;

            auto median = std::make_shared<ScalarAggregateFunction>(
                    "approximate_median", Arity::Unary(), approximate_median_doc,
                    &default_scalar_aggregate_options);

            auto sig = KernelSignature::create({InputType::Any()}, float64());

            auto init = [tdigest_func](
                    KernelContext *ctx,
                    const KernelInitArgs &args) -> turbo::Result<std::unique_ptr<KernelState>> {
                std::vector<TypeHolder> types = args.inputs;
                TURBO_MOVE_OR_RAISE(auto kernel, tdigest_func->dispatch_best(&types));
                const auto &scalar_options =
                        turbo::checked_cast<const ScalarAggregateOptions &>(*args.options);
                TDigestOptions options;
                // Default q = 0.5
                options.min_count = scalar_options.min_count;
                options.skip_nulls = scalar_options.skip_nulls;
                KernelInitArgs new_args{kernel, types, &options};
                return kernel->init(ctx, new_args);
            };

            auto finalize = [](KernelContext *ctx, Datum *out) -> turbo::Status {
                Datum temp;
                TURBO_RETURN_NOT_OK(turbo::checked_cast<ScalarAggregator *>(ctx->state())->Finalize(ctx, &temp));
                const auto arr = temp.make_array();
                DKCHECK_EQ(arr->length(), 1);
                return arr->get_scalar(0).try_value(out);
            };

            add_agg_kernel(std::move(sig), std::move(init), std::move(finalize), median.get());
            return median;
        }

    }  // namespace

    void register_scalar_aggregate_tdigest(FunctionRegistry *registry) {
        auto tdigest = add_tdigest_agg_kernels();
        KCHECK_OK(registry->add_function(tdigest));

        auto approx_median = add_approximate_median_agg_kernels(tdigest.get());
        KCHECK_OK(registry->add_function(approx_median));
    }

}  // namespace nebula::compute::internal

